# -*- coding: utf-8 -*-

# from __future__ import print_function
import os
import random
# import sys
import numpy as np
import cv2
import argparse
# import PIL
from PIL import ImageFont
from PIL import Image
from PIL import ImageDraw
import math

font_ch_plate = 'cigeratte_tf/fonts/drumnarrow.ttf'

plate_template_img_path = 'cigeratte_tf/bg/4.jpg'

plate_noise_img_path = 'plate_tf/images/smu2.jpg'

plate_height_default = 128
plate_width_default = plate_height_default * 3

cigeratte_chars_count = 32

# INDEX_DIGIT = {"0": 31, "1": 32, "2": 33, "3": 34, "4": 35, "5": 36, "6": 37, "7": 38, "8": 39, "9": 40}
#
# INDEX_LETTER = {"A": 41, "B": 42, "C": 43, "D": 44, "E": 45, "F": 46, "G": 47,
#                 "H": 48, "J": 49, "K": 50, "L": 51, "M": 52, "N": 53,
#                 "P": 54, "Q": 55, "R": 56, "S": 57, "T": 58,
#                 "U": 59, "V": 60, "W": 61, "X": 62, "Y": 63, "Z": 64}

PLATE_CHARS_DIGIT = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]

PLATE_CHARS_LETTER = ["A", "B", "C", "D", "E", "F", "G",
                        "H", "I", "J", "K", "L", "M", "N",
                        "O", "P", "Q", "R", "S", "T",
                        "U", "V", "W", "X", "Y", "Z"]

# PLATE_INDEX_DICT = dict(INDEX_DIGIT, **INDEX_LETTER)

PLATE_CHARS_TOTAL = PLATE_CHARS_DIGIT + PLATE_CHARS_LETTER

# print PLATE_INDEX_DICT,PLATE_CHARS_TOTAL

def append_noise(img, Smu):
    img_h, img_w = img.shape[:2]
    rows = rand_int(Smu.shape[0] - img_h)

    cols = rand_int(Smu.shape[1] - img_w)
    adder = Smu[rows:rows + img_h, cols:cols + img_w]
    adder = cv2.resize(adder, (img_w, img_h))
    adder = cv2.bitwise_not(adder)
    val = random.random() * 0.3
    img = cv2.addWeighted(img, 1 - val, adder, val, 0.0)
    return img

def rot(img,angel,max_angel):
    """ 使图像轻微的畸变
        img 输入图像
        factor 畸变的参数
        size 为图片的目标尺寸
    """
    shape = img.shape
    size_o = [shape[1],shape[0]]
    # print size_o
    # size = (shape[1]+ int(shape[0]*cos((float(max_angel )/180) * 3.14)),shape[0])
    # print size
    size = (shape[1]+ int(shape[0] * math.sin((float(max_angel )/180) * 3.14)),shape[0])
    # print size
    interval = abs( int( math.sin((float(angel) /180) * 3.14)* shape[0]))

    pts1 = np.float32([[0,0],[0,size_o[1]],[size_o[0],0],[size_o[0],size_o[1]]])
    if(angel>0):
        pts2 = np.float32([[interval,0],[0,size[1]  ],[size[0],0  ],[size[0]-interval,size_o[1]]])
    else:
        pts2 = np.float32([[0,0],[interval,size[1]  ],[size[0]-interval,0  ],[size[0],size_o[1]]])

    M  = cv2.getPerspectiveTransform(pts1,pts2)
    dst = cv2.warpPerspective(img,M,size)

    return dst

def rotRandrom(img, width_offset, height_offset):
    size = (img.shape[1],img.shape[0])
    # pts1 = np.float32([[0, 0], [0, size[0]], [size[1], 0], [size[1], size[0]]])
    # pts2 = np.float32([[rand_int(height_offset), rand_int(width_offset)],
    #                    [rand_int(height_offset), size[0] - rand_int(width_offset)],
    #                    [size[1] - rand_int(height_offset), rand_int(width_offset)],
    #                    [size[1] - rand_int(height_offset), size[0] - rand_int(width_offset)]])
    pts1 = np.float32([[0, 0], [size[0],0], [0,size[1]], [size[0],size[1]]])
    pts2 = np.float32([[rand_int(width_offset), rand_int(height_offset)],
                       [size[0] - rand_int(width_offset),rand_int(height_offset) ],
                       [rand_int(width_offset),size[1] - rand_int(height_offset)],
                       [size[0] - rand_int(width_offset),size[1] - rand_int(height_offset)]])
    # 仿射不行，需要投影
    # affine = cv2.getAffineTransform(pts1,pts2)
    # dst = cv2.warpAffine(img,affine,size)
    M = cv2.getPerspectiveTransform(pts1, pts2)
    dst = cv2.warpPerspective(img, M, size)
    return dst

def tfactor(img):
    hsv = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)

    hsv[:,:,0] = hsv[:,:,0]*(0.8+ np.random.random()*0.2)
    hsv[:,:,1] = hsv[:,:,1]*(0.3+ np.random.random()*0.7)
    hsv[:,:,2] = hsv[:,:,2]*(0.2+ np.random.random()*0.8)

    img = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)
    return img

def random_environment(img, data_set):
    index=rand_int(len(data_set))
    env = cv2.imread(data_set[index])

    env = cv2.resize(env,(img.shape[1],img.shape[0]))

    # bak = (img==0);
    # bak = bak.astype(np.uint8)*255;
    # inv = cv2.bitwise_and(bak,env)
    # img = cv2.bitwise_or(inv,img)
    val = random.random() * 0.4
    img = cv2.addWeighted(img, 1 - val, env, val, 0.0)
    return img

def random_scene(img, data_set):
    '''将车牌放入自然场景图像中，并返回该图像和位置信息'''
    bg_img_path = data_set[rand_int(len(data_set))]
    # print bg_img_path
    env = cv2.imread(bg_img_path)
    if env is None:
        print bg_img_path, 'is not a good file'
        return None, None
    # print env.shape, img.shape
    # 如果背景图片小于（65，21）则不使用
    if env.shape[1] <= 65 or env.shape[0] <= 21:
        print env.shape
        return None, None
    # 车牌宽高比变化范围是(1.5, 4.0)
    new_height = img.shape[0] * (0.5 + np.random.random()) # 0.5 -- 1.5
    new_width = img.shape[1] * (0.5 + np.random.random()) # 0.5 -- 1.5
    scale = 0.3 + np.random.random() * 2.5
    new_width = int(new_width * scale + 0.5)
    new_height = int(new_height * scale + 0.5)
    img = cv2.resize(img, (new_width, new_height))
    if env.shape[1] <= img.shape[1] or env.shape[0] <= img.shape[0]:
        print env.shape, '---', img.shape
        return None, None
    x = rand_int(env.shape[1] - img.shape[1])
    y = rand_int(env.shape[0] - img.shape[0])
    bak = (img==0)
    bak = bak.astype(np.uint8)*255
    inv = cv2.bitwise_and(bak, env[y:y+new_height, x:x+new_width, :])
    img = cv2.bitwise_or(inv, img)
    env[y:y+new_height, x:x+new_width, :] = img[:,:,:]

    return env, (x, y, x + img.shape[1], y + img.shape[0])

def draw_text(f, img, text):
    # img=Image.new("RGB", (23,70),(255,255,255))
    pilimage = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pilimage)
    draw.text((0, 10), text[:16].decode('utf-8'), (255, 255, 255), font=f)
    draw.text((0, 46), text[16:].decode('utf-8'), (255, 255, 255), font=f)
    return cv2.cvtColor(np.asarray(pilimage), cv2.COLOR_RGB2BGR)

def draw_text_2(f, text):
    width = 506
    height = 140
    pilimage = Image.new("RGB",(width,height))
    draw = ImageDraw.Draw(pilimage)
    # 274 * 40
    size1 = draw.textsize(text[:16].decode('utf-8'),f)
    size2 = draw.textsize(text[16:].decode('utf-8'),f)
    clr = rand_int(220,180)
    draw.text((0, 0), text[:16].decode('utf-8'), (clr,clr,clr), font=f)
    draw.text((0, 70 ), text[16:].decode('utf-8'), (clr,clr,clr), font=f)
    pilimage = cv2.cvtColor(np.asarray(pilimage), cv2.COLOR_RGB2BGR)
    # cv2.imshow("img",pilimage)
    # cv2.waitKey(0)
    return pilimage

def draw_text_3(f, text):
    width = 324
    height = 90
    pilimage = Image.new("RGB",(width,height))
    draw = ImageDraw.Draw(pilimage)
    # 274 * 40
    size1 = draw.textsize(text[:16].decode('utf-8'),f)
    size2 = draw.textsize(text[16:].decode('utf-8'),f)
    clr = rand_int(230,180)
    draw.text((0, 0), text[:16].decode('utf-8'), (clr,clr,clr), font=f)
    draw.text((0, 45 ), text[16:].decode('utf-8'), (clr,clr,clr), font=f)
    pilimage = cv2.cvtColor(np.asarray(pilimage), cv2.COLOR_RGB2BGR)
    return pilimage

def rotate_img(img,degree,scale = 1.0):
    height,width=img.shape[:2]
    height_new = width * math.fabs(math.sin(math.radians(degree))) + height * math.fabs(math.cos(math.radians(degree)))
    width_new = height * math.fabs(math.sin(math.radians(degree))) + width * math.fabs(math.cos(math.radians(degree)))
    height_new = int(height_new * scale)
    width_new = int(width_new * scale)
    matRotation=cv2.getRotationMatrix2D((width/2,height/2),degree,scale)
    matRotation[0, 2] += (width_new - width) / 2
    matRotation[1, 2] += (height_new - height) / 2
    dst = cv2.warpAffine(img, matRotation, (width_new, height_new))
    # cv2.imshow("dst",dst)
    # cv2.waitKey(0)
    return dst


# 低通线性滤波
def add_blur(img, level):
    return cv2.blur(img, (level * 2 + 1, level * 2 + 1))

# 高斯滤波
def add_gauss_blur(img, level):
    return cv2.GaussianBlur(img, (level * 2 + 1, level * 2 + 1), 1.5)

def rand_int(max,min = 0):
    return int(np.random.random() * (max-min))+min

def add_noise_single_channel(single):
    diff = 255-single.max()
    noise = np.random.normal(0, 1 + rand_int(6), single.shape)
    noise = (noise - noise.min())/(noise.max()-noise.min())
    noise= diff*noise
    noise= noise.astype(np.uint8)
    dst = single + noise
    return dst

def add_noise(img, sdev = 0.5, avg=10):
    img[:,:,0] =  add_noise_single_channel(img[:, :, 0])
    img[:,:,1] =  add_noise_single_channel(img[:, :, 1])
    img[:,:,2] =  add_noise_single_channel(img[:, :, 2])
    return img

# 腐蚀
def Erosion(src, es):
    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * es + 1, 2 * es + 1), (es, es))
    return cv2.erode( src, element )

# 膨胀
def Dilation(src, es):
    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * es + 1, 2 * es + 1), (es, es))
    return cv2.dilate( src, element )
